{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 导入库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import math\n",
    "import random\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 生成数据"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "+ **数据格式说明**  \n",
    "  数据集的格式为{graph1, graph2, ..., graphn}\n",
    "  其中，graphi为第i个graph数据，graphi的数据结构为G_i(N, E)：\n",
    "        N = {v_1, v_2, ..., v_n}  # n个节点集合\n",
    "        E = {e_1, e_2, ..., e_t}  # t条边集合\n",
    "   其中，e_j的数据结构为(v_p, v_q)，表示该边两端的节点。\n",
    "+ **任务说明**  \n",
    "  给定子图sub_graph，训练集{graph1, graph2, ..., graphn}，模型目标就是对其中的graph中的节点进行分类。  \n",
    "  如果graphi包含sub_graph和节点v，且节点v属于sub_graph中的一个节点，那么节点v的标签为1，反之，标签为-1。模型就是要预测每个节点的标签是1还是-1。\n",
    "+ **其它说明**  \n",
    "  这里使用的数据直接按照论文《The Graph Neural Network Model》所述的子图匹配任务进行实现，节点的特征就是数字1-10  \n",
    "    1, 2, 3, 4, 5, 6, 7, 8, 9, 10  \n",
    "  因此，节点集合使用节点特征的列表表示，边集合中的每一条边(i,j)中的i和j表示该边两端的节点在节点列表中对应的索引。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "使用networkx库生成graph的图像，并显示，节点中显示的数字表示节点特征数字。\n",
    "Input :\n",
    "    graph : (Node_list, Edge_list)\n",
    "'''\n",
    "def graphFigure(graph):\n",
    "    G = nx.Graph()\n",
    "    node_feat_list, edge_list = graph\n",
    "    node_names = [\"v{}\".format(i) for i in range(len(node_feat_list))]\n",
    "    G.add_nodes_from(node_names, feat=node_feat_list)\n",
    "    G.add_edges_from([(\"v{}\".format(i),\"v{}\".format(j)) for i,j in edge_list])\n",
    "    labels = dict([(\"v{}\".format(i), feat) for i, feat in enumerate(node_feat_list)])\n",
    "    plt.figure(0)\n",
    "    nx.draw_spectral(G=G, \n",
    "                     node_size=1000,\n",
    "                     node_color=\"g\",\n",
    "                     with_labels=True, \n",
    "                     font_weight='bold',\n",
    "                     labels=labels)\n",
    "'''\n",
    "用于产生一个graph数据，首先随机生成若干个节点，然后随机对这些节点\n",
    "进行连接，得到一个随机的graph，然后将subgraph插入产生的随机graph，\n",
    "得到graph数据。\n",
    "Input :\n",
    "    sub_graph : 子图数据\n",
    "    node_num : 生成随机graph的节点个数\n",
    "    edge_num : 生成随机graph的边条数\n",
    "    ins_edge_num : 将subgraph插入随机graph时，连接到随机graph的边条数\n",
    "'''\n",
    "def genGraph(sub_graph, node_num, edge_num, ins_edge_num):\n",
    "    nodes_list = list(np.random.randint(low=1, high=11, size=(node_num)))\n",
    "    label_list = [-1] * len(nodes_list)\n",
    "    edges_list = []\n",
    "    edge_proba = edge_num / (node_num * (node_num - 1) / 2)\n",
    "    for n in range(node_num-1):\n",
    "        end_nodes = np.random.choice(a=np.arange(n+1, node_num), \n",
    "                                     size=(round(edge_proba * (node_num-1-n))),\n",
    "                                     replace=False)\n",
    "        edges_list.extend([(n, e) for e in end_nodes])\n",
    "    # Insert subgraph\n",
    "    nodes_list.extend(sub_graph[0])  # add subgraph nodes\n",
    "    label_list.extend([1]*len(sub_graph[0]))\n",
    "    head_nodes = np.random.choice(a=np.arange(node_num),\n",
    "                                  size=(ins_edge_num))\n",
    "    tail_nodes = np.random.choice(a=np.arange(node_num, len(nodes_list)),\n",
    "                                  size=(ins_edge_num))\n",
    "    edges_list.extend([(i,j) for i,j in zip(head_nodes, tail_nodes)])\n",
    "    edges_list.extend([(i+node_num,j+node_num) for i,j in sub_graph[1]])\n",
    "    return (nodes_list, edges_list, label_list)\n",
    "    \n",
    "'''\n",
    "Definition of subgraph and dataset generation\n",
    "'''\n",
    "np.random.seed(0)\n",
    "N = 600          # graph dataset length\n",
    "node_nums = [5, 10, 15, 20]\n",
    "ins_nums = [4, 8, 12, 16]\n",
    "\n",
    "\n",
    "sub_graph = ([1,5,5,8],[(0,1),(0,3),(1,3),(2,3)])  # (Node_list, Edge_list)\n",
    "dataset = []\n",
    "for i in range(N):\n",
    "    nnum = random.choice(node_nums)   # node_num\n",
    "    ins_en = random.choice(ins_nums)# ins_edge_num\n",
    "    dataset.append(genGraph(sub_graph, nnum, 2*nnum, ins_en))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 构建模型"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GCN构建"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "同Linear GNN\n",
    "'''\n",
    "class AggrSum(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(AggrSum, self).__init__()\n",
    "    \n",
    "    def forward(self, H, X_node, node_num):\n",
    "        # H : (N, s) -> (V, s)\n",
    "        # X_node : (N, )\n",
    "        mask = torch.stack([X_node] * node_num, 0)\n",
    "        mask = mask.float() - torch.unsqueeze(torch.range(0,node_num-1).float(), 1)\n",
    "        mask = (mask == 0).float()\n",
    "        # (V, N) * (N, s) -> (V, s)\n",
    "        return torch.mm(mask, H)\n",
    "\n",
    "'''\n",
    "用于实现GCN的卷积块。\n",
    "Initialize :\n",
    "Input :\n",
    "    in_channel : (int)输入的节点特征维度\n",
    "    out_channel : (int)输出的节点特征维度\n",
    "Forward :\n",
    "Input :\n",
    "    x : (Tensor)节点的特征矩阵，shape为(N, in_channel)，N为节点个数\n",
    "    edge_index : (Tensor)边矩阵，shape为(2, E)，E为边个数。\n",
    "Output :\n",
    "    out : (Tensor)新的特征矩阵，shape为(N, out_channel)\n",
    "'''\n",
    "class GCNConv(nn.Module):\n",
    "    def __init__(self, in_channel, out_channel):\n",
    "        super(GCNConv, self).__init__()\n",
    "        self.linear = nn.Linear(in_channel, out_channel)\n",
    "        self.aggregation = AggrSum()\n",
    "        \n",
    "    def forward(self, x, edge_index):\n",
    "        # Add self-connect edges\n",
    "        edge_index = self.addSelfConnect(edge_index, x.shape[0])\n",
    "        node_num = x.shape[0]\n",
    "        \n",
    "        # Apply linear transform\n",
    "        x = self.linear(x)\n",
    "        \n",
    "        # Normalize message\n",
    "        row, col = edge_index\n",
    "        deg = self.calDegree(row, x.shape[0]).float()\n",
    "        deg_sqrt = deg.pow(-0.5)  # (N, )\n",
    "        norm = deg_sqrt[row] * deg_sqrt[col]\n",
    "        \n",
    "        # Node feature matrix\n",
    "        tar_matrix = torch.index_select(x, dim=0, index=col)\n",
    "        tar_matrix = norm.view(-1, 1) * tar_matrix  # (E, out_channel)\n",
    "        # Aggregate information\n",
    "        aggr =  self.aggregation(tar_matrix, row, node_num)  # (N, out_channel)\n",
    "        return aggr\n",
    "        \n",
    "    def calDegree(self, edges, num_nodes):\n",
    "        ind, deg = np.unique(edges.cpu().numpy(), return_counts=True)\n",
    "        deg_tensor = torch.zeros((num_nodes, ), dtype=torch.long)\n",
    "        deg_tensor[ind] = torch.from_numpy(deg)\n",
    "        return deg_tensor.to(edges.device)\n",
    "    \n",
    "    def addSelfConnect(self, edge_index, num_nodes):\n",
    "        selfconn = torch.stack([torch.range(0, num_nodes-1, dtype=torch.long)]*2,\n",
    "                               dim=0).to(edge_index.device)\n",
    "        return torch.cat(tensors=[edge_index, selfconn],\n",
    "                         dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 两层GCN构建"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "构建模型，使用两层GCN，第一层GCN使得节点特征矩阵\n",
    "    (N, in_channel) -> (N, out_channel)\n",
    "第二层GCN直接输出\n",
    "    (N, out_channel) -> (N, num_class)\n",
    "激活函数使用relu函数，网络最后对节点的各个类别score使用softmax归一化。\n",
    "'''\n",
    "class Net(nn.Module):\n",
    "    def __init__(self, feat_dim, num_class):\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = GCNConv(feat_dim, 16)\n",
    "        self.conv2 = GCNConv(16, num_class)\n",
    "    \n",
    "    def forward(self, x, edge_index):\n",
    "        x = self.conv1(x, edge_index)\n",
    "        x = F.relu(x)\n",
    "        x = F.dropout(x, training=self.training)\n",
    "        x = self.conv2(x, edge_index)\n",
    "        \n",
    "        return F.softmax(x, dim=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 准备输入数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding = np.diag(np.ones((10)))\n",
    "\n",
    "'''\n",
    "得到x和edge_index输入\n",
    "'''\n",
    "def getInput(graph):\n",
    "    x = embedding[[i-1 for i in graph[0]]]\n",
    "    edge_index = np.array([np.array([i,j]) for i,j in graph[1]]).T  # (2, E)\n",
    "    edge_index = np.concatenate([edge_index] * 2, axis=1)   # (2, 2*E)\n",
    "    y = np.array(graph[2])\n",
    "    return x, edge_index, y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 开始训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "开始训练模型\n",
    "'''\n",
    "device = torch.device('cpu')# torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model = Net(10, 2).to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)\n",
    "# loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "'''\n",
    "模型评估函数\n",
    "'''\n",
    "def evalModel(model, dataset):\n",
    "    for graph in dataset:\n",
    "        x, edge_index, y = getInput(graph)\n",
    "        x = torch.from_numpy(x).float()\n",
    "        edge_index = torch.from_numpy(edge_index).long()\n",
    "        y[y < 0] = 0\n",
    "        y = torch.from_numpy(y).long()\n",
    "        \n",
    "        acc_list = []\n",
    "        _, pred = model(x, edge_index).max(dim=1)\n",
    "        acc_list.append(float(pred.eq(y).sum().item())/y.shape[0])\n",
    "    return sum(acc_list)/ len(acc_list)\n",
    "\n",
    "for epoch in range(200):\n",
    "    for step, graph in enumerate(dataset[:400]):\n",
    "        # Get input\n",
    "        x, edge_index, y = getInput(graph)\n",
    "        x = torch.from_numpy(x).float().to(device)\n",
    "        edge_index = torch.from_numpy(edge_index).long().to(device)\n",
    "        y[y < 0] = 0\n",
    "        y = torch.from_numpy(y).long().to(device)\n",
    "        \n",
    "        model.train()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Get output\n",
    "        out = model(x, edge_index)   # (N, 2)\n",
    "        # Get loss\n",
    "        loss = F.cross_entropy(out, y)\n",
    "\n",
    "        # Backward\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        # Get predictions and calculate training accuracy\n",
    "        _, pred = out.cpu().detach().max(dim=-1)  # (N)\n",
    "        y = y.cpu().detach()\n",
    "        correct = float(pred.eq(y).sum().item())\n",
    "        acc = correct / pred.shape[0]\n",
    "        print('[Epoch {}/200, step {}/400] Loss {:.4f}, train acc {:.4f}'.format(epoch, step, loss.cpu().detach().data.item(), acc))\n",
    "\n",
    "    # Evaluation on test data every 10 epochs\n",
    "    if (epoch+1) % 10 == 0:\n",
    "        model.eval()\n",
    "        print('Accuracy: {:.4f}'.format(evalModel(model, dataset[400:])))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
